using System.Runtime.InteropServices;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;
using Random = UnityEngine.Random;

namespace Learning
{
    public class Simple3GPUInstancing : MonoBehaviour
    {
        /// <summary>
        /// GPU Instancing 一个批次可提交的数量
        /// 注意：这里设置得越大，DC会越少，但是特殊情况下浪费的带宽会越多（比如只提交10个时，一个批次还是会提交512个数据给GPU，但是有效数据只有10个）
        /// </summary>
        private const int BATCH = 512;
        
        private const float Range = 100000f;
        private readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));
        
        public int instanceCount = 100000;
        public Mesh instanceMesh;
        public Material instanceMaterial;

        private TransInfo[] _transInfos; // CPU 计算的每个对象的信息
        
        void Start()
        {
            CreateInfo();
        }

        private void CreateInfo()
        {
            _transInfos = new TransInfo[instanceCount];
            for (int i = 0; i < instanceCount; i++)
            {
                var info = new TransInfo();
                info.pos = new float3(Random.Range(-1000f, 1000f), Random.Range(-1000f, 1000f), Random.Range(-1000f, 1000f));
                info.rotate = quaternion.Euler(Random.Range(0f, 360f), Random.Range(0f, 360f), Random.Range(0f, 360f));
                info.scale = new float3(Random.Range(1f, 5f), Random.Range(1f, 5f), Random.Range(1f, 5f));
                info.forward = new float3(Random.Range(0f, 1f), Random.Range(0f, 1f), Random.Range(0f, 1f));
                info.rotateAdd = new float3(Random.Range(0f, 10f), Random.Range(0f, 10f), Random.Range(0f, 10f));
                _transInfos[i] = info;
            }
        }

        void Update()
        {
            UpdateInfo();

            var batchCount = (int) math.ceil((float) instanceCount / (float) BATCH);
            for (int i = 0; i < batchCount; i++)
            {
                var startIndex = i * BATCH;
                var len = BATCH;
                if (i == batchCount - 1)
                {
                    len = instanceCount - (BATCH * i);
                }

                // 这里应该做成对象池，这里图方便就暴力创建了
                var batchInfo = new BatchInfo();
                
                // 填充数据，数据长度固定才能使用对象池
                var transArr = new float4x4[BATCH];
                for (int j = 0; j < len; j++)
                {
                    var info = _transInfos[startIndex + j];
                    transArr[j] = float4x4.TRS(info.pos, info.rotate, info.scale);
                }
                var transBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(float4x4)));
                transBuf.SetData(transArr);
                batchInfo.block = new MaterialPropertyBlock();
                batchInfo.block.SetBuffer("transBuf", transBuf);
                
                var args = new uint[5] {0, 0, 0, 0, 0};
                args[0] = instanceMesh.GetIndexCount(0); // 每个实例的索引数
                args[1] = (uint) len; // 实例数
                args[2] = instanceMesh.GetIndexStart(0); // 起始索引位置
                args[3] = instanceMesh.GetBaseVertex(0); // 基顶点位置
                args[4] = 0; // 起始实例位置
                batchInfo.argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
                batchInfo.argsBuf.SetData(args);

                // 绘制
                Graphics.DrawMeshInstancedIndirect(instanceMesh, 0, instanceMaterial, BOUNDS, batchInfo.argsBuf,
                    0, batchInfo.block, ShadowCastingMode.On, true, LayerMask.NameToLayer("Character"));
            }
        }

        private void UpdateInfo()
        {
            for (int i = 0; i < instanceCount; i++)
            {
                var info = _transInfos[i];
                info.pos += info.forward * Time.deltaTime;
                info.rotate = math.mul(info.rotate, quaternion.Euler(info.rotateAdd * Time.deltaTime));
            }
        }
    }

    public class BatchInfo
    {
        public MaterialPropertyBlock block;
        public GraphicsBuffer argsBuf;
    }
}
